import os
import re
import jieba
import ast
import numpy as np
import pandas as pd
import imageio
import matplotlib.pyplot as plt
from wordcloud import WordCloud

# -------------------------- 1. 配置基础参数 --------------------------
# 路径配置（请根据你的文件位置修改！）
DATA_DIR = "../data"  # 存放positive/negative_samples.txt、停用词、背景图的目录
SAVE_DIR = "../tmp"   # 词云图保存目录
BACKGROUND_PIC = os.path.join(DATA_DIR, "background.jpg")  # 词云背景图（建议用酒店/风景相关图片）
STOPWORD_FILE = os.path.join(DATA_DIR, "stopword.txt")    # 停用词文件
CUSTOM_DICT = os.path.join(DATA_DIR, "newdic1.txt")        # 自定义词典（可加酒店领域词如“大床房”“自助餐”）

# 确保保存目录存在
if not os.path.exists(SAVE_DIR):
    os.makedirs(SAVE_DIR)

# -------------------------- 2. 读取并解析酒店评论TXT --------------------------
def parse_hotel_txt(file_path, sentiment_label):
    """
    解析酒店评论TXT（每行是{'text': '评论内容', 'label': 1/0}格式）
    :param file_path: TXT文件路径
    :param sentiment_label: 情感标签（正面/负面）
    :return: 包含“评论内容+情感标签”的DataFrame
    """
    comment_list = []
    with open(file_path, 'r', encoding='utf-8') as f:
        for line_num, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue  # 跳过空行
            try:
                # 解析字典字符串（处理可能的格式错误）
                comment_dict = ast.literal_eval(line)
                # 提取评论内容，过滤过短评论（小于10字的可能是无效内容）
                if len(comment_dict.get('text', '')) >= 10:
                    comment_list.append({
                        'comment': comment_dict['text'],
                        'sentiment': sentiment_label
                    })
            except Exception as e:
                print(f"第{line_num}行格式错误，已跳过：{e}")
    return pd.DataFrame(comment_list)

# 读取正面/负面评论（根据你的negative_samples.txt标签调整，此处假设label=0为负面）
positive_df = parse_hotel_txt(
    file_path=os.path.join(DATA_DIR, "positive_samples.txt"),
    sentiment_label="正面"
)
negative_df = parse_hotel_txt(
    file_path=os.path.join(DATA_DIR, "negative_samples.txt"),
    sentiment_label="负面"
)

# 合并数据并查看基本信息
hotel_data = pd.concat([positive_df, negative_df], ignore_index=True)
print(f"总评论数：{len(hotel_data)}")
print(f"情感分布：\n{hotel_data['sentiment'].value_counts()}")

# -------------------------- 3. 酒店评论预处理（针对酒店场景优化） --------------------------
def preprocess_comment(comment_series):
    """
    评论预处理：去重、脱敏、分词、去停用词
    :param comment_series: 评论内容的Series
    :return: 预处理后的分词列表（无停用词）
    """
    # 3.1 去重
    comment_dup = comment_series.drop_duplicates()
    print(f"去重后评论数：{len(comment_dup)}")

    # 3.2 脱敏（去除评论中的“x”“***”等隐私信息）
    comment_qumin = comment_dup.astype(str).apply(
        lambda x: re.sub(r'x+|\*+', '', x)  # 匹配多个x或*
    )

    # 3.3 加载自定义词典（补充酒店领域词汇，如“双床房”“行政楼层”）
    if os.path.exists(CUSTOM_DICT):
        jieba.load_userdict(CUSTOM_DICT)
        print("已加载酒店领域自定义词典")

    # 3.4 分词（保留长度≥2的词，过滤单字无意义词）
    comment_cut = comment_qumin.apply(
        lambda x: [word for word in jieba.cut(x) if len(word) >= 2]
    )

    # 3.5 去停用词（新增酒店领域通用词如“酒店”“房间”，避免占比过高）
    # 基础停用词
    stopwords = pd.read_csv(STOPWORD_FILE, sep='ooo', encoding='gbk', header=None, engine='python')[0].tolist()
    # 酒店领域停用词
    hotel_stopwords = [
        '酒店', '房间', '入住', '客人', '我们', '这里', '什么', '没有', '一个', '这个', '那个',
        '携程', '事情', '情况', '发现', '建议', '问题', '告诉', '2008', '10', '25', '2007',
        '人员', '同事', '朋友', '大家', '们', '者', '些', '方面', '时候', '时间', '地方',
        '东西', '方式', '部分', '认为', '觉得', '前台', '宾馆', '设施','服务员','环境','只能',
        '分钟','总体','电梯','地点','小姐','洗澡','设计','对面','品种','每次',
        '打电话', '告知', '通知', '电梯', '大堂', '总台', '窗户', '12', '20',  # 场景/数字类
        '因为', '所以', '虽然', '但是', '而且', '或者', '不过',  # 连接词类
        '能够', '可以', '应该', '需要', '想要', '得', '要', '会',  # 助动词类
        '感觉', '知道', '看到', '听到', '觉得', '认为', '得到', '做到',  # 感知/判断类
        '出来', '进去', '里面', '外面', '左右', '上下', '前后', '之间',  # 方位词类
        '非常', '比较', '太', '很', '挺',  # 程度副词
    ]
    all_stopwords = stopwords + hotel_stopwords + [' ']  # 合并停用词

    # 过滤停用词
    comment_qustop = comment_cut.apply(
        lambda x: [word for word in x if word not in all_stopwords]
    )

    # 3.6 删除空列表（分词后无有效词的评论）
    comment_final = comment_qustop[comment_qustop.astype(str) != '[]']
    print(f"预处理后有效评论数：{len(comment_final)}")
    return comment_final

# 执行预处理
processed_comments = preprocess_comment(hotel_data['comment'])

# 关联情感标签（确保预处理后评论与标签对应）
comment_labels = hotel_data.loc[processed_comments.index, 'sentiment']

# -------------------------- 4. 酒店评论词频统计 --------------------------
def count_word_freq(comment_list, min_freq=15):
    """
    统计高频词（过滤低频词，避免词云杂乱）
    :param comment_list: 预处理后的分词列表
    :param min_freq: 最小出现频率（酒店评论建议15-20）
    :return: 高频词Series（词：频率）
    """
    # 合并所有评论的词
    all_words = []
    for words in comment_list:
        all_words.extend(words)
    # 统计词频
    word_freq = pd.Series(all_words).value_counts()
    # 保留高频词
    high_freq_words = word_freq[word_freq >= min_freq]
    print(f"高频词数量：{len(high_freq_words)}")
    print("前10个高频词：\n", high_freq_words.head(10))
    return high_freq_words

# 按情感分类统计高频词
positive_comments = processed_comments[comment_labels == "正面"]
negative_comments = processed_comments[comment_labels == "负面"]

positive_freq = count_word_freq(positive_comments, min_freq=15)  # 正面评论高频词
negative_freq = count_word_freq(negative_comments, min_freq=10)  # 负面评论样本可能少，降低频率阈值

# -------------------------- 5. 生成酒店评论词云图 --------------------------
def generate_wordcloud(high_freq_words, save_name, background_pic):
    """
    生成词云图并保存
    :param high_freq_words: 高频词Series
    :param save_name: 保存文件名（如“酒店正面评论词云”）
    :param background_pic: 背景图路径
    """
    # 配置词云参数（适配中文，优化视觉效果）
    wc = WordCloud(
        font_path='C:/Windows/Fonts/simkai.ttf',  # 中文黑体字体（确保路径存在）
        background_color='white',                  # 背景色（白色显干净）
        max_words=1500,                           # 最大词数（避免拥挤）
        mask=imageio.imread(background_pic),       # 背景图轮廓
        max_font_size=180,                        # 最大字体大小
        random_state=42,                          # 固定配色方案
        collocations=False                         # 避免重复词组（如“服务好”“好服务”）
    )

    # 生成词云
    wc.generate_from_frequencies(high_freq_words)

    # 显示并保存
    plt.figure(figsize=(12, 8))  # 调整画布大小（适配背景图）
    plt.imshow(wc, interpolation='bilinear')  # 平滑显示
    plt.axis('off')  # 隐藏坐标轴
    plt.title(
        f'{save_name}',
        fontsize=16,
        pad=20,
        fontproperties='SimHei')  # 标题
    # 保存图片（高清格式）
    save_path = os.path.join(SAVE_DIR, f"{save_name}.png")
    plt.savefig(save_path, dpi=300, bbox_inches='tight')  # dpi=300保证高清
    plt.show()
    print(f"{save_name}已保存至：{save_path}")

# 生成正面/负面评论词云
if not positive_freq.empty:
    generate_wordcloud(positive_freq, "酒店正面评论词云", BACKGROUND_PIC)
else:
    print("正面评论高频词不足，无法生成词云")

if not negative_freq.empty:
    generate_wordcloud(negative_freq, "酒店负面评论词云", BACKGROUND_PIC)
else:
    print("负面评论高频词不足，无法生成词云")